import threading
import os.path
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QMessageBox, QMainWindow, QApplication
from util import *
import MainWindow
from drone_by_robomaster import drone_by_robomaster
from drone_by_simulation import simulation_drone
from yoloDetections_to_npArray import tran
import yolov5

pressed_key_lock = threading.Lock()
debug = True
debug = False


def thread_it(func, *args):
    t = threading.Thread(target=func, args=args)
    t.setDaemon(True)
    t.start()
    return t


class My_MainWindow(QMainWindow, MainWindow.Ui_MainWindow):
    def __init__(self, parent=None):
        super(My_MainWindow, self).__init__(parent)
        self.setupUi(self)
        # 核心组件
        self.drone = None

        # 标志位
        self.exit_status = False  # 是否退出程序
        self.camera_status = True  # 是否打开相机
        self.take_off_status = False  # 是否已经起飞
        self.auto_crossing_ring_status = False  # 是否进入自主运行（穿越）模式
        self.auto_tracking_balloon_status = False  # 是否进入自主运行（追踪气球）模式

        self.auto_crossing_ring_pose_up = False  # 自主运行模式（穿越）中，是否摆正姿势
        self.auto_crossing_ring_sprint = False  # 自主运行模式（穿越）中，是否开始执行穿越
        self.auto_crossing_ring_sprint_begin_time = None  # 自主运行模式（穿越）中，开始穿越的时间

        # 内置变量
        self.pressed_keys = set()
        self.lastest_img = None  # 无人机图传最新帧
        self.img = None  # 最新展示帧
        self.image_recognition_results = None  # 图像识别结果（来自self.img）
        # {
        #   0: 'UAV',
        #   1: 'UGV',
        #   2: 'balloon',
        #   3: 'block',
        #   4: 'box',
        #   5: 'cirque',
        #   6: 'duck',
        #   7: 'person'
        # }

        # 渲染窗口
        self.setWindowTitle('无人机控制')
        self.setWindowIcon(QtGui.QIcon('favicon.ico'))
        self.setFixedSize(self.width(), self.height())
        self.pushButton_3.setEnabled(False)
        self.pushButton_5.setEnabled(False)
        self.statusBar.showMessage("正在初始化...", 0)
        self.show()
        # self.model = yolov5.load('yolov5s.pt')
        self.model = yolov5.load('balloon_recognition_model.pt')

        # 连接无人机
        success_connect = False
        log('正在连接无人机...')
        try:
            if not debug:
                self.drone = drone_by_robomaster()
            else:
                self.drone = simulation_drone()
            success_connect = True
            log('连接成功')
        except:
            QMessageBox.critical(self, '错误', '连接无人机失败')
            self.close_service()

        # 启动服务线程
        if success_connect:
            self.init()

    def init(self):
        thread_it(self.event_loop)
        thread_it(self.open_camera)
        thread_it(self.show_qt_video)
        self.pushButton.clicked.connect(lambda: thread_it(self.switch_camera_status))
        self.pushButton_2.clicked.connect(lambda: thread_it(self.switch_takeoff_switch))
        self.pushButton_3.clicked.connect(lambda: thread_it(self.switch_auto_crossing_ring_switch))
        self.pushButton_5.clicked.connect(lambda: thread_it(self.switch_auto_tracking_balloon_switch))
        self.pushButton_4.clicked.connect(lambda: thread_it(self.emergency_stop))
        self.lcdNumber.display(self.drone.drone_get_gear())

    def show_cv_img(self):
        shrink = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
        QtImg = QtGui.QImage(shrink.data,
                             shrink.shape[1],
                             shrink.shape[0],
                             shrink.shape[1] * 3,
                             QtGui.QImage.Format_RGB888)
        jpg_out = QtGui.QPixmap(QtImg).scaled(
            self.label.width(), self.label.height())
        self.label.setPixmap(jpg_out)

    def open_camera(self):
        self.statusBar.showMessage("执行指令 " + "打开相机", 0)
        self.drone.drone_init_camera()
        self.statusBar.showMessage("就绪 " + "打开相机", 0)
        self.camera_status = True
        if not os.path.exists('imgs'):
            os.mkdir('imgs')
        while self.camera_status and not self.exit_status:
            self.lastest_img = self.drone.drone_get_img()

    def process_qt_img(self, img):
        image_recognition_results = self.model(img)
        self.image_recognition_results = image_recognition_results

        img = tran(image_recognition_results)

        if self.auto_crossing_ring_status:
            pred_width, pred_height, pred_center_x, pred_center_y, img_center_x, img_center_y, threshold_x, threshold_y, dx, dy = self.auto_crossing_ring_evaluating_positional_deviation(
                img, image_recognition_results)
        if self.auto_tracking_balloon_status:
            pred_width, pred_height, pred_center_x, pred_center_y, img_center_x, img_center_y, threshold_x, threshold_y, dx, dy = self.auto_tracking_balloon_evaluating_positional_deviation(
                img, image_recognition_results)

        if self.auto_crossing_ring_status or self.auto_tracking_balloon_status:
            if pred_width:

                # 在img上画出阈值框
                # 如果需要调整就画红色的框，否则画绿色的框
                color = (0, 0, 255) if dx or dy else (0, 255, 0)
                cv2.rectangle(img, (int(img_center_x - threshold_x), int(img_center_y - threshold_y)),
                              (int(img_center_x + threshold_x), int(img_center_y + threshold_y)), color, 2)
                # 画出预测框的中心点，点的颜色为蓝色，半径为5
                cv2.circle(img, (int(pred_center_x), int(pred_center_y)), 5, (255, 0, 0), -1)
                # 画一条从图像中心点到预测框中心点的线，线的颜色为蓝色，线的粗细为2
                cv2.line(img, (int(img_center_x), int(img_center_y)), (int(pred_center_x), int(pred_center_y)),
                         (255, 0, 0), 2)
            else:
                color = (0, 0, 255)
                cv2.rectangle(img, (int(img_center_x - threshold_x), int(img_center_y - threshold_y)),
                              (int(img_center_x + threshold_x), int(img_center_y + threshold_y)), color, 2)
        return img

    def show_qt_video(self):
        while not self.exit_status:
            if self.lastest_img is None or not self.camera_status:
                time.sleep(0.01)
                continue
            # beg_time = time.time()
            # end_time = time.time()
            # print('model time: ', end_time - beg_time)
            self.img = self.process_qt_img(self.lastest_img)
            self.show_cv_img()
            # cv2.waitKey(1)

    def switch_takeoff_switch(self):
        self.drone.drone_stop()
        if self.take_off_status:
            self.take_off_status = False
            self.pushButton_2.setText("起飞")
            self.pushButton_2.setEnabled(False)
            self.pushButton_3.setEnabled(False)
            self.pushButton_5.setEnabled(False)
            self.drone.drone_land()
            self.pushButton_2.setEnabled(True)
        else:
            self.take_off_status = True
            self.pushButton_2.setText("降落")
            self.pushButton_2.setEnabled(False)
            self.drone.drone_takeoff()
            self.pushButton_2.setEnabled(True)
            self.pushButton_3.setEnabled(True)
            self.pushButton_5.setEnabled(True)

    def switch_auto_crossing_ring_switch(self):
        # 目标标签为：cirque
        self.target_label = 5
        if self.auto_crossing_ring_status:
            self.auto_crossing_ring_status = False
            self.pushButton_3.setText("穿越圆环")

            # 启用其他自动化功能
            self.pushButton_5.setEnabled(True)
            # 允许关闭相机
            self.pushButton.setEnabled(True)

            self.drone.drone_stop()
        else:
            self.auto_crossing_ring_status = True
            self.auto_crossing_ring_pose_up = False
            self.auto_crossing_ring_sprint = False

            # 禁用其他自动化功能
            self.pushButton_5.setEnabled(False)
            # 禁止关闭相机
            self.pushButton.setEnabled(False)

            self.pushButton_3.setText("停止运行")

    def switch_auto_tracking_balloon_switch(self):
        # 目标标签为：balloon
        self.target_label = 2
        if self.auto_tracking_balloon_status:
            self.auto_tracking_balloon_status = False
            self.pushButton_5.setText("追踪气球")

            # 启用其他自动化功能
            self.pushButton_3.setEnabled(True)
            # 允许关闭相机
            self.pushButton.setEnabled(True)

            self.drone.drone_stop()
        else:
            self.auto_tracking_balloon_status = True

            # 禁用其他自动化功能
            self.pushButton_3.setEnabled(False)
            # 禁止关闭相机
            self.pushButton.setEnabled(False)

            self.pushButton_5.setText("停止运行")

    def switch_camera_status(self):
        if self.camera_status:
            self.statusBar.showMessage("执行指令 " + "关闭相机", 0)
            self.camera_status = False
            self.drone.drone_close_camera()
            self.statusBar.showMessage("就绪 " + "关闭相机", 0)
            self.pushButton.setText("打开相机")
        else:
            self.pushButton.setText("关闭相机")
            thread_it(self.open_camera)

    def emergency_stop(self):
        if self.auto_crossing_ring_status:
            self.switch_auto_crossing_ring_switch()
        if self.auto_tracking_balloon_status:
            self.switch_auto_tracking_balloon_switch()
        if self.take_off_status:
            self.switch_takeoff_switch()
        else:
            self.drone.drone_land()

    # 监控键盘
    def keyPressEvent(self, event):
        key = event.key()
        # log('keyPressEvent', key)

        with pressed_key_lock:
            if key == Qt.Key_W:
                self.pressed_keys.add('forward')
            elif key == Qt.Key_S:
                self.pressed_keys.add('backward')
            elif key == Qt.Key_A:
                self.pressed_keys.add('left')
            elif key == Qt.Key_D:
                self.pressed_keys.add('right')

            elif key == Qt.Key_I:
                self.pressed_keys.add('up')
            elif key == Qt.Key_J:
                self.pressed_keys.add('left_rotate')
            elif key == Qt.Key_K:
                self.pressed_keys.add('down')
            elif key == Qt.Key_L:
                self.pressed_keys.add('right_rotate')
            elif key == Qt.Key_Return:
                self.pressed_keys.add('flip')

            elif key == Qt.Key_Plus or key == Qt.Key_U:
                self.pressed_keys.add('plus')
            elif key == Qt.Key_Minus or key == Qt.Key_H:
                self.pressed_keys.add('minus')

            elif key == Qt.Key_N:
                self.pressed_keys.add('takeoff')
            elif key == Qt.Key_M:
                self.pressed_keys.add('land')
            elif key == Qt.Key_Shift:
                self.pressed_keys.add('accelerate')

    def keyReleaseEvent(self, event):
        key = event.key()
        # log('keyReleaseEvent', key)
        with pressed_key_lock:
            if key == Qt.Key_W:
                self.pressed_keys.remove('forward')
            elif key == Qt.Key_S:
                self.pressed_keys.remove('backward')
            elif key == Qt.Key_A:
                self.pressed_keys.remove('left')
            elif key == Qt.Key_D:
                self.pressed_keys.remove('right')

            elif key == Qt.Key_I:
                self.pressed_keys.remove('up')
            elif key == Qt.Key_J:
                self.pressed_keys.remove('left_rotate')
            elif key == Qt.Key_K:
                self.pressed_keys.remove('down')
            elif key == Qt.Key_L:
                self.pressed_keys.remove('right_rotate')
            elif key == Qt.Key_Return:
                self.pressed_keys.remove('flip')

            elif key == Qt.Key_Plus or key == Qt.Key_U:
                self.pressed_keys.remove('plus')
            elif key == Qt.Key_Minus or key == Qt.Key_H:
                self.pressed_keys.remove('minus')

            elif key == Qt.Key_N:
                self.pressed_keys.remove('takeoff')
            elif key == Qt.Key_M:
                self.pressed_keys.remove('land')
            elif key == Qt.Key_Shift:
                self.pressed_keys.remove('accelerate')

    def auto_crossing_ring_evaluating_positional_deviation(self, img, image_recognition_results):
        pred_list = image_recognition_results.pred[0].tolist()
        pred_list = [i[:5] + [int(i[5])] for i in pred_list]

        pred_target_list = []
        for pred in pred_list:
            if int(pred[5]) == self.target_label:
                pred_target_list.append(pred)

        # 求图像的中心点
        img_center_y, img_center_x = img.shape[0] * 0.4, img.shape[1] / 2

        # 设定一个阈值，当差值小于阈值时，认为无需调整
        threshold_x = 25
        threshold_y = 20

        if len(pred_target_list) == 0:
            return None, None, None, None, img_center_x, img_center_y, threshold_x, threshold_y, None, None

        # 选出距离中心点最近的识别结果
        min_distance = 100000000
        min_distance_index = 0
        for i, pred in enumerate(pred_target_list):
            distance = (pred[0] - img_center_x) ** 2 + (pred[1] - img_center_y) ** 2
            if distance < min_distance:
                min_distance = distance
                min_distance_index = i
        pred_target = pred_target_list[min_distance_index]

        # 计算预测框的长和宽
        pred_width = pred_target[2] - pred_target[0]
        pred_height = pred_target[3] - pred_target[1]

        # 计算预测框的中心点
        pred_center_x = (pred_target[0] + pred_target[2]) / 2
        pred_center_y = (pred_target[1] + pred_target[3]) / 2

        # 分别计算预测框的中心点和图像中心点的x和y的差值
        dx = pred_center_x - img_center_x
        dy = pred_center_y - img_center_y

        # 判断是否需要调整
        if abs(dx) < threshold_x:
            dx = 0
        if abs(dy) < threshold_y:
            dy = 0

        return pred_width, pred_height, pred_center_x, pred_center_y, img_center_x, img_center_y, threshold_x, threshold_y, dx, dy

    def auto_tracking_balloon_evaluating_positional_deviation(self, img, image_recognition_results):
        pred_list = image_recognition_results.pred[0].tolist()
        pred_list = [i[:5] + [int(i[5])] for i in pred_list]

        # {0: 'UAV', 1: 'UGV', 2: 'balloon', 3: 'block', 4: 'box', 5: 'cirque', 6: 'duck', 7: 'person'}
        # 目标标签为：balloon
        target_label = 2
        pred_target_list = []
        for pred in pred_list:
            if int(pred[5]) == target_label:
                pred_target_list.append(pred)

        # 求图像的中心点
        img_center_y, img_center_x = img.shape[0] / 2, img.shape[1] / 2

        # 设定一个阈值，当差值小于阈值时，认为无需调整
        threshold_x = 90
        threshold_y = 90

        if len(pred_target_list) == 0:
            return None, None, None, None, img_center_x, img_center_y, threshold_x, threshold_y, None, None

        # 选出距离中心点最近的识别结果
        min_distance = 100000000
        min_distance_index = 0
        for i, pred in enumerate(pred_target_list):
            distance = (pred[0] - img_center_x) ** 2 + (pred[1] - img_center_y) ** 2
            if distance < min_distance:
                min_distance = distance
                min_distance_index = i
        pred_target = pred_target_list[min_distance_index]

        # 计算预测框的长和宽
        pred_width = pred_target[2] - pred_target[0]
        pred_height = pred_target[3] - pred_target[1]

        # 计算预测框的中心点
        pred_center_x = (pred_target[0] + pred_target[2]) / 2
        pred_center_y = (pred_target[1] + pred_target[3]) / 2

        # 分别计算预测框的中心点和图像中心点的x和y的差值
        dx = pred_center_x - img_center_x
        dy = pred_center_y - img_center_y

        # 判断是否需要调整
        if abs(dx) < threshold_x:
            dx = 0
        if abs(dy) < threshold_y:
            dy = 0

        return pred_width, pred_height, pred_center_x, pred_center_y, img_center_x, img_center_y, threshold_x, threshold_y, dx, dy

    def drone_auto_crossing_ring(self, img, image_recognition_results):
        pred_width, pred_height, pred_center_x, pred_center_y, img_center_x, img_center_y, threshold_x, threshold_y, dx, dy = self.auto_crossing_ring_evaluating_positional_deviation(
            img, image_recognition_results)
        if pred_width and not self.auto_crossing_ring_sprint:
            auto_crossing_ring_pose_up = True

            a, b, c, d, speed_magnification = 0, 0, 0, 0, 1
            if dx > 0 and not self.auto_crossing_ring_sprint:  # 若正在穿越，则不进行左右调整
                # 向右调整
                d = -1
                auto_crossing_ring_pose_up = False
            elif dx < 0 and not self.auto_crossing_ring_sprint:  # 若正在穿越，则不进行左右调整
                # 向左调整
                d = 1
                auto_crossing_ring_pose_up = False

            if dy > 0:
                # 向上调整
                c = -1
                auto_crossing_ring_pose_up = False
            elif dy < 0:
                # 向下调整
                c = 1
                auto_crossing_ring_pose_up = False

            if pred_width < pred_height and \
                    abs(pred_width * 0.6 - pred_height) / pred_width > 0.2 and \
                    not self.auto_crossing_ring_sprint:  # 若正在穿越，则不进行前右移动
                # 向右移动
                a = -1
                auto_crossing_ring_pose_up = False

            # 判断角度是否修正完毕，若是，则开始接近目标
            if auto_crossing_ring_pose_up:
                self.auto_crossing_ring_pose_up = True
            if self.auto_crossing_ring_pose_up:
                speed_magnification = 2
                b = 1

            # 判断是否接近目标，若是，则开始穿越
            if self.auto_crossing_ring_pose_up and pred_width > 0.7 * img.shape[
                1] and not self.auto_crossing_ring_sprint:
                self.auto_crossing_ring_sprint = True
                self.auto_crossing_ring_sprint_begin_time = time.time()
                log('开始穿越')

            self.drone.drone_move(a, b, c, d, speed_magnification)

        if self.auto_crossing_ring_sprint:
            a, b, c, d, speed_magnification = 0, 0, 0, 0, 2
            b = 1

            if pred_width:
                if dy > 0:
                    # 向上调整
                    c = -1
                elif dy < 0:
                    # 向下调整
                    c = 1

            # 退出穿越
            if self.auto_crossing_ring_sprint and time.time() - self.auto_crossing_ring_sprint_begin_time > 1.75:
                log('退出穿越，花费时间：%.2f' % (time.time() - self.auto_crossing_ring_sprint_begin_time))
                self.switch_auto_crossing_ring_switch()
                self.auto_crossing_ring_pose_up = False
                self.auto_crossing_ring_sprint = False

            self.drone.drone_move(a, b, c, d, speed_magnification)

        if not pred_width and not self.auto_crossing_ring_sprint:
            self.drone.drone_stop()

        # 显示当前任务执行状态
        if not self.auto_crossing_ring_pose_up:
            self.statusBar.showMessage("正在调整无人机的角度", 0)
        elif not self.auto_crossing_ring_sprint:
            self.statusBar.showMessage("正在接近目标", 0)
        elif pred_width:
            self.statusBar.showMessage("正在穿越", 0)
        else:
            self.statusBar.showMessage("未检测到目标，待命", 0)

    def drone_auto_tracking_balloons(self, img, image_recognition_results):
        pred_width, pred_height, pred_center_x, pred_center_y, img_center_x, img_center_y, threshold_x, threshold_y, dx, dy = self.auto_tracking_balloon_evaluating_positional_deviation(img, image_recognition_results)
        if pred_width:
            a, b, c, d, speed_magnification = 0, 0, 0, 0, 4
            if dx > 0:
                # 向右调整
                d = -1
            elif dx < 0:
                # 向左调整
                d = 1

            if dy > 0:
                # 向上调整
                c = -1
            elif dy < 0:
                # 向下调整
                c = 1
            self.drone.drone_move(a, b, c, d, speed_magnification)

    def event_loop(self):
        last_pressed_keys = set()
        while not self.exit_status:
            if self.take_off_status and self.auto_crossing_ring_status:
                self.statusBar.showMessage('正在执行穿越圆环', 0)
                self.drone_auto_crossing_ring(self.img, self.image_recognition_results)

            if self.take_off_status and self.auto_tracking_balloon_status:
                self.statusBar.showMessage('正在执行追踪气球', 0)
                self.drone_auto_tracking_balloons(self.img, self.image_recognition_results)


            with pressed_key_lock:
                for key in self.pressed_keys:
                    if key not in last_pressed_keys:
                        log('按下了' + key, 'keys', self.pressed_keys)
                for key in last_pressed_keys:
                    if key not in self.pressed_keys:
                        log('松开了' + key, 'keys', self.pressed_keys)

            if self.pressed_keys != last_pressed_keys:
                last_pressed_keys = self.pressed_keys.copy()
                status_text = ''
                if len(last_pressed_keys) == 0:
                    status_text = '停止移动'
                    self.statusBar.showMessage('执行指令 ' + status_text, 0)
                    self.drone.drone_stop()
                    self.statusBar.showMessage('就绪 ' + status_text, 0)
                else:
                    if 'takeoff' in last_pressed_keys:
                        status_text = "起飞"
                        self.statusBar.showMessage('执行指令 ' + status_text, 0)
                        if not self.take_off_status:
                            self.switch_takeoff_switch()
                        self.statusBar.showMessage("就绪 " + status_text, 0)
                    elif 'land' in last_pressed_keys:
                        status_text = "降落"
                        self.statusBar.showMessage('执行指令 ' + status_text, 0)
                        if self.take_off_status:
                            self.switch_takeoff_switch()
                        self.statusBar.showMessage("就绪 " + status_text, 0)
                    elif 'flip' in last_pressed_keys:
                        direction = ''
                        if 'forward' in last_pressed_keys:
                            direction = 'f'
                        elif 'backward' in last_pressed_keys:
                            direction = 'b'
                        elif 'left' in last_pressed_keys:
                            direction = 'l'
                        elif 'right' in last_pressed_keys:
                            direction = 'r'
                        if direction != '':
                            status_text = "翻转"
                        self.statusBar.showMessage('执行指令 ' + status_text, 0)
                        self.drone.drone_flip(direction)
                        self.statusBar.showMessage("就绪 " + status_text, 0)
                    else:
                        a, b, c, d, speed_magnification = 0, 0, 0, 0, 1
                        for key in last_pressed_keys:
                            if key == 'forward':
                                b = 1
                                status_text += '前进 '
                            elif key == 'backward':
                                b = -1
                                status_text += '后退 '
                            elif key == 'left':
                                a = 1
                                status_text += '左移 '
                            elif key == 'right':
                                a = -1
                                status_text += '右移 '
                            elif key == 'up':
                                c = 1
                                status_text += '上升 '
                            elif key == 'down':
                                c = -1
                                status_text += '下降 '
                            elif key == 'left_rotate':
                                d = 1
                                status_text += '左转 '
                            elif key == 'right_rotate':
                                d = -1
                                status_text += '右转 '
                            elif key == 'plus':
                                speed_magnification += 1
                                status_text += '加速 '
                                self.drone.drone_increase_gear()
                                self.lcdNumber.display(self.drone.drone_get_gear())
                            elif key == 'minus':
                                speed_magnification -= 1
                                self.drone.drone_decrease_gear()
                                self.lcdNumber.display(self.drone.drone_get_gear())
                                status_text += '减速 '
                            elif key == 'accelerate':
                                speed_magnification *= 2
                                status_text += '紧急加速 '
                        self.statusBar.showMessage('执行指令 ' + status_text, 0)
                        self.drone.drone_move(a, b, c, d, speed_magnification)
                        self.statusBar.showMessage("就绪 " + status_text, 0)
            self.lcdNumber_2.display(self.drone.drone_get_battery())
            time.sleep(0.001)

    def close_service(self):
        self.exit_status = True
        try:
            self.drone.drone_stop()
            self.drone.drone_land()
            self.drone.close()
        except:
            pass

    # 退出
    def closeEvent(self, event):
        self.close_service()


if __name__ == '__main__':
    import sys

    try:
        app = QApplication(sys.argv)
        ui = My_MainWindow()
        sys.exit(app.exec_())
    except RuntimeError:
        pass
